import torch  # 导入PyTorch核心库
import matplotlib.pyplot as plt  # 用于绘制损失曲线
import numpy as np  # 数值计算库（本文件未直接使用，可能为后续扩展保留）
import torch.utils.data as data  # 数据加载工具
import config  # 导入项目配置参数
import torchvision  # 计算机视觉工具库（本文件未直接使用，可能为后续扩展保留）
from net import YoloV1Net  # 导入YOLOv1网络模型
from dataset import YoloV1Dataset  # 导入自定义数据集类
from loss import YoloV1Loss  # 导入YOLOv1损失函数


# ---------------------------
# 数据加载与预处理
# ---------------------------
print ("Loading the datasets")  # 打印数据加载状态
# 初始化完整数据集（启用归一化）
data_tv = YoloV1Dataset ("./datasets/train.csv", "./datasets/images", "./datasets/labels", normalize=True)
# 划分训练集和验证集（8:2比例）
train_data, val_data = data.random_split(data_tv, [0.8, 0.2])

# 训练数据加载器（批量加载+随机打乱）
train_loader = data.DataLoader (
    dataset=train_data,
    batch_size=config.BATCH_SIZE,  # 批量大小（config中定义）
    shuffle=True  # 每轮训练打乱数据顺序
)
# 验证数据加载器（批量加载+随机打乱）
val_loader = data.DataLoader (
    dataset=val_data,
    batch_size=config.BATCH_SIZE,
    shuffle=True
)


# ---------------------------
# 模型初始化与配置
# ---------------------------
model = YoloV1Net ()  # 初始化YOLOv1模型
model.to ('cuda')  # 将模型移动到GPU（若可用）

print (model)  # 打印模型结构（检查网络是否正确）

optimizer = torch.optim.Adam (model.parameters(), lr=config.LR)  # 使用Adam优化器（学习率config中定义）
loss_func = YoloV1Loss ()  # 初始化YOLOv1损失函数


# ---------------------------
# 训练前信息打印
# ---------------------------
print (f"Epochs : {config.EPOCHS}\nTrain_set Length : {train_data.__len__ ()}\nBatch Size : {config.BATCH_SIZE}")
print ("Start Training ...")  # 打印训练开始提示


# ---------------------------
# 主训练循环
# ---------------------------
for epoch in range (config.EPOCHS) :  # 遍历所有训练轮次
    train_loss = []  # 记录每10步的训练损失
    for step, (input, label) in enumerate (train_loader) :  # 遍历训练数据加载器
        input, label = input.to ('cuda'), label.to ('cuda')  # 将数据移动到GPU
        out = model (input)  # 前向传播（获取模型预测）
        loss = loss_func (out, label)  # 计算预测与标签的损失

        optimizer.zero_grad ()  # 清空优化器梯度
        loss.backward ()  # 反向传播（计算梯度）
        optimizer.step()  # 更新模型参数

        # 每10个批次打印一次损失并记录
        if step % 10 == 0 :
            print (f"Epoch : {epoch}\tLoss : {loss.item ()}")  # 打印当前轮次和损失
            train_loss.append(loss.item ())  # 记录损失值用于后续绘图


# ---------------------------
# 训练结果保存与可视化
# ---------------------------
plt.plot (train_loss)  # 绘制训练损失曲线
plt.show ()  # 显示损失曲线
torch.save(model, "./model.pt")  # 保存训练好的模型（路径：当前目录下的model.pt）
